from multiprocessing.dummy.connection import Connection

import cv2
import torch

from cjc import project_setting
from cjc.detector import opt_tool
from cjc.detector.detecotor_single import DetectorSingle
from cjc.detector.opt_tool import create_opt


class AgeDetector(DetectorSingle):
    def __init__(self, _opt):
        _opt.weights = project_setting.root+'/cjc/weights/age.pt'
        _opt.source = '-1'
        _opt.name = 'age'
        super(AgeDetector, self).__init__(_opt)


def processor_run_age_dector(p: Connection, q):
    opt = opt_tool.create_opt()
    det = AgeDetector(opt)
    classes = det.names
    while True:
        cv2.waitKey(1)
        img_set = p.recv()
        if img_set is not None:
            a = det.detect(img_set.img, str(img_set.id))
            if a:
                a = a[0].split(' ')
                msg = [img_set.id,
                       'age',
                       a[-1],
                       classes[int(a[0])]
                       ]
                q.put(msg)
                print("放入了一个数据", msg)


if __name__ == '__main__':
    opt = create_opt()
    opt.view_img = True
    detector = AgeDetector(opt)
    # f='/mnt/d/datasets/sex_and_age/sex/train/images/'
    # file='39_1_4_20170105001146812.jpg'
    # f1=f'{f}{file}'
    # frame= cv2.imread(f1)
    # a = detector.detect(frame, 0)
    # cv2.waitKey(10000)
    # print(a)

    cv_cap = cv2.VideoCapture(2)
    while True:
        ret, frame = cv_cap.read()
        if ret:
            with torch.no_grad():
                a = detector.detect(frame, 0)
                cv2.waitKey(1)
                print(a)